# import model_faction as model
# import datasources as data

import LinReg.model_faction as model
import LinReg.datasources as data

import tensorflow as tf

lr = 0.03  # 学习率
num_epochs = 3  # 训练凑
batch_size = 10  # 批量大小

# 初始化w, b
w = tf.Variable(tf.random.normal(shape=(2, 1), mean=0, stddev=0.01), trainable=True)
b = tf.Variable(tf.zeros(1), trainable=True)

"""
这段代码使用梯度下降算法训练一个线性回归模型。它接受多个参数，包括训练轮数 num_epochs，批量大小 batch_size，学习率 lr，数据集 data，模型参数 w 和 b 以及模型对象 model。

首先，它使用一个外层 for 循环遍历每个训练轮。在每个训练轮中，它使用一个内层 for 循环遍历数据集中的每个批次。

对于每个批次，它使用 tf.GradientTape 记录梯度信息，并计算损失函数 l。然后，它使用 g.gradient 函数计算模型参数的梯度，并使用 model.sgd 函数更新模型参数。

在每个训练轮结束时，它计算并打印出训练损失和模型参数的估计误差。

这段代码的目的是使用梯度下降算法训练一个线性回归模型，并在每个训练轮结束时输出训练进度。
"""
for epoch in range(num_epochs):
    for X, y in model.data_iter(batch_size, data.sample, data.labels):
        with tf.GradientTape() as g:
            l = model.squared_loss(model.linreg(X, w, b), y)
        dw, db = g.gradient(l, [w, b])
        model.sgd([w, b], [dw, db], lr, batch_size)
    train_l = model.squared_loss(model.linreg(data.sample, w, b), data.labels)
    print(f'epoch{epoch + 1}, loss{float(tf.reduce_mean(train_l)): f}')
    print(f'w的估计误差：{data.true_w - tf.reshape(w, data.true_w.shape)}')
    print(f'b的估计误差{data.true_b - b}')
